#include <opencv2/opencv.hpp>
#include <iostream>
#include <chrono>
#include <X11/Xlib.h>
#include <X11/Xutil.h>
#include <X11/extensions/XShm.h>
#include <sys/shm.h>

#include <vector>
#include <thread>
#include <queue>
#include <mutex>
#include <sys/socket.h>
#include <netinet/in.h>
#include <unistd.h>

#include "tcpserver.hpp"
#include "semaphore.hpp"
#include "imageProcess.hpp"

#define Height 800
#define Width 1280
#define HeightTransmit 600
#define WidthTransmit 960
#define SegmentWidth (64)
#define SegmentHeight (40)
#define FrameInterval 125

enum class ThreadState
{
    Idle = 0,
    Running = 1,
    Exit = 2,
};

typedef struct
{
    uint32_t left;
    uint32_t top;
    uint32_t width;
    uint32_t height;
} Setting;

int PORT = 6117;
int Quality = 80;

static bool is_running = True;
static Display *display;
static Window root;
static XImage *ximage;

static Semaphore sem_grab;
static Semaphore sem_send;

static std::mutex mtx_frame_buffer;
static std::mutex mtx_data_stream;
static std::thread compress_thread;
static ThreadState compress_thread_is_running = ThreadState::Idle;

static std::vector<ImageBlock> image_blocks;
static std::vector<uchar> data_stream;

#pragma pack(push, 8) // 设置结构体按8字节对齐
static uchar frame_buffer[Height * Width * 4];
static uchar frame_buffer_transmit[HeightTransmit * WidthTransmit * 4];
static uchar frame_buffer_copy[HeightTransmit * WidthTransmit * 4];
#pragma pack(pop) // 恢复默认对齐方式

static void compressThread();
static int handleFunction(int client_fd, std::string cmd);

static int handleFunction(int client_fd, std::string cmd)
{
    std::vector<uchar> _data_stream;

    uint32_t left;
    uint32_t top;
    uint32_t width;
    uint32_t height;
    uint32_t jpeg_len;

    try
    {
        Quality = std::stoi(cmd);

        if (compress_thread_is_running == ThreadState::Running)
        {
            std::cout << "compress_thread_is_running, wait for thread exiting" << "\n";
            compress_thread_is_running = ThreadState::Exit;
            compress_thread.join();
            std::cout << "compress_thread has been exited" << "\n";
        }

        compress_thread = std::thread(compressThread);

        while (is_running)
        {
            auto start = std::chrono::high_resolution_clock::now();

            sem_send.wait_for_signal();

            mtx_data_stream.lock();
            std::vector<uchar> _data_stream = data_stream;
            mtx_data_stream.unlock();

            uint32_t len = _data_stream.size();
            if (send(client_fd, (uchar *)&len, sizeof(uint32_t), MSG_NOSIGNAL) < 0) // MSG_NOSIGNAL: 避免出现 Broken pipe
                return 0;

            // if (len == 0)
            // {
            //     std::this_thread::sleep_for(std::chrono::milliseconds(20));
            //     continue;
            // }

            if (send(client_fd, _data_stream.data(), _data_stream.size(), MSG_NOSIGNAL) < 0) // MSG_NOSIGNAL: 避免出现 Broken pipe
                return 0;

            auto end = std::chrono::high_resolution_clock::now();
            std::chrono::duration<double, std::milli> elapsed = end - start;
            std::cout << "preprocess image Elapsed time: " << elapsed.count() << " ms\n";

            std::cout << "\n";
        }
    }
    catch (const std::exception &e)
    {
        std::cerr << "Exception: " << e.what() << std::endl;
        return 0;
    }

    // close(client_fd);
    return 1;
}

static void compressThread()
{
    uint32_t left;
    uint32_t top;
    uint32_t width;
    uint32_t height;
    uint32_t jpeg_len;
    std::vector<uchar> _data_stream;

    compress_thread_is_running = ThreadState::Running;

    memset(frame_buffer_copy, 0xff, sizeof(frame_buffer_copy));
    auto frame_time_start = std::chrono::high_resolution_clock::now();

    // std::this_thread::sleep_for(std::chrono::seconds(1));
    while (compress_thread_is_running == ThreadState::Running)
    {
        auto start = std::chrono::high_resolution_clock::now();

        mtx_frame_buffer.lock();
        // memcopy(frame_buffer_copy, frame_buffer, Height * Width * sizeof(uint32_t));
        imageBlockingCompareAndCopy(frame_buffer_transmit, frame_buffer_copy, WidthTransmit, HeightTransmit, SegmentWidth, SegmentHeight, image_blocks);
        cv::Mat image(HeightTransmit, WidthTransmit, CV_8UC4, (char *)frame_buffer_copy);
        mtx_frame_buffer.unlock();

        auto end = std::chrono::high_resolution_clock::now();
        std::chrono::duration<double, std::milli> elapsed = end - start;
        std::cout << "preprocess image Elapsed time: " << elapsed.count() << " ms\n";

        sem_grab.send_signal();
        auto _start = std::chrono::high_resolution_clock::now();

        /******** 原图发送 ***********/
        // std::vector<uchar> __data_stream(frame_buffer_transmit, frame_buffer_transmit + WidthTransmit * HeightTransmit * 4);
        // mtx_data_stream.lock();
        // data_stream = __data_stream;
        // mtx_data_stream.unlock();
        // sem_send.send_signal();
        // continue;

        _data_stream.clear();
        /********************** 分块压缩 ************************/
        for (std::vector<ImageBlock>::iterator it = image_blocks.begin(); it != image_blocks.end(); ++it)
        {
            // 定义矩形区域 (x, y, Width, height) 截取图像
            cv::Rect rect(it->left, it->top, it->width, it->height);
            cv::Mat roi = image(rect);

            // 编码图像为 JPEG 字节流
            std::vector<uchar> _jpeg_buffer;
            _jpeg_buffer.clear();
            cv::imencode(".jpg", roi, _jpeg_buffer, std::vector<int>({cv::IMWRITE_JPEG_QUALITY, Quality}));
            // std::cout << "jpg " << it->left << " " << it->top << " data len: " << _jpeg_buffer.size() << " \n";
            jpeg_len = _jpeg_buffer.size();
            _data_stream.insert(_data_stream.end(), reinterpret_cast<uchar *>(&*it), reinterpret_cast<uchar *>(&*it) + sizeof(ImageBlock));
            _data_stream.insert(_data_stream.end(), (uchar *)&jpeg_len, (uchar *)&jpeg_len + sizeof(jpeg_len));
            _data_stream.insert(_data_stream.end(), _jpeg_buffer.begin(), _jpeg_buffer.end());
        }

        std::cout << "image_blocks.size: " << image_blocks.size() << "\n";

        /********************** 整图压缩 ************************/
        // std::vector<uchar> jpeg_buffer;
        // cv::imencode(".jpg", image, jpeg_buffer, std::vector<int>({cv::IMWRITE_JPEG_QUALITY, Quality}));

        // jpeg_len = jpeg_buffer.size();
        // ImageBlock ib = {0, 0, Width, Height};
        // _data_stream.insert(_data_stream.end(), reinterpret_cast<uchar *>(&ib), reinterpret_cast<uchar *>(&ib) + sizeof(ImageBlock));
        // _data_stream.insert(_data_stream.end(), (uchar *)&jpeg_len, (uchar *)&jpeg_len + sizeof(jpeg_len));
        // _data_stream.insert(_data_stream.end(), jpeg_buffer.begin(), jpeg_buffer.end());
        /**********************************************/

        auto _end = std::chrono::high_resolution_clock::now();
        std::chrono::duration<double, std::milli> _elapsed = _end - _start;
        std::cout << "imencode block jpeg Elapsed time: " << _elapsed.count() << " ms\n";

        mtx_data_stream.lock();
        data_stream = _data_stream;
        mtx_data_stream.unlock();

        auto frame_time_end = std::chrono::high_resolution_clock::now();
        std::chrono::duration<double, std::milli> frame_time_elapsed = frame_time_end - frame_time_start;
        int frame_time_elapsed_int_value = static_cast<int>(frame_time_elapsed.count());
        auto time_ms = FrameInterval - frame_time_elapsed_int_value;
        if (time_ms > 0)
        // if (image_blocks.size() == 0)
        {
            std::this_thread::sleep_for(std::chrono::milliseconds(time_ms));
        }
        frame_time_start = std::chrono::high_resolution_clock::now();

        sem_send.send_signal();
    }
    compress_thread_is_running = ThreadState::Idle;
}

void grabService()
{
    TCPServer server(PORT, handleFunction);
    server.start();
    std::cout << "Server is running on port " << PORT << "..." << std::endl;

    XShmSegmentInfo shminfo;
    display = XOpenDisplay(NULL);
    root = DefaultRootWindow(display);

    XWindowAttributes attributes;
    XGetWindowAttributes(display, root, &attributes);

    // 创建共享内存图像
    ximage = XShmCreateImage(display,
                             DefaultVisual(display, DefaultScreen(display)),
                             attributes.depth,
                             ZPixmap,
                             NULL,
                             &shminfo,
                             attributes.width, attributes.height);

    // XShmSegmentInfo shminfo;
    shminfo.shmid = shmget(IPC_PRIVATE, ximage->bytes_per_line * ximage->height, IPC_CREAT | 0777);
    shminfo.shmaddr = ximage->data = (char *)shmat(shminfo.shmid, 0, 0);
    shminfo.readOnly = False;

    XShmAttach(display, &shminfo);

    while (is_running)
    {
        sem_grab.wait_for_signal();

        XSync(display, False);
        XShmGetImage(display, root, ximage, 0, 0, AllPlanes);

        if (Height != ximage->height || Width != ximage->width)
        {
            std::cout << "Screen resolution error: " << ximage->width << "x" << ximage->height << "\n";
            continue;
        }

        auto start = std::chrono::high_resolution_clock::now();

        if (Width == WidthTransmit && Height == HeightTransmit)
        {
            mtx_frame_buffer.lock();
            memcopy(frame_buffer_transmit, (uchar *)ximage->data, ximage->bytes_per_line * ximage->height);
            mtx_frame_buffer.unlock();
        }
        else
        {
            memcopy(frame_buffer, (uchar *)ximage->data, ximage->bytes_per_line * ximage->height);
            mtx_frame_buffer.lock();
            resize(frame_buffer, frame_buffer_transmit, Width, Height, WidthTransmit, HeightTransmit);
            mtx_frame_buffer.unlock();
        }

        auto end = std::chrono::high_resolution_clock::now();
        std::chrono::duration<double, std::milli> elapsed = end - start;
        std::cout << "XShmGetImage Elapsed time: " << elapsed.count() << " ms\n";
    }

    shmdt(shminfo.shmaddr);
    XShmDetach(display, &shminfo);
    XDestroyImage(ximage);
    XCloseDisplay(display);

    server.stop();
    std::cout << "Server stopped." << std::endl;
}
